#include "EncryptionHelper.h"
#include <openssl/evp.h>
#include <openssl/rand.h>
#include <openssl/bio.h>
#include <openssl/buffer.h>
#include <sstream>
#include <iomanip>
#include <stdexcept>
#include <cstring>

namespace MessageFramework {

// ==================== Base64编解码 ====================

std::string EncryptionHelper::base64Encode(const std::vector<unsigned char>& data) {
    BIO* bio = BIO_new(BIO_s_mem());
    BIO* b64 = BIO_new(BIO_f_base64());
    bio = BIO_push(b64, bio);
    BIO_set_flags(bio, BIO_FLAGS_BASE64_NO_NL);  // 不换行
    BIO_write(bio, data.data(), static_cast<int>(data.size()));
    BIO_flush(bio);

    BUF_MEM* bufferPtr;
    BIO_get_mem_ptr(bio, &bufferPtr);
    std::string result(bufferPtr->data, bufferPtr->length);

    BIO_free_all(bio);
    return result;
}

std::vector<unsigned char> EncryptionHelper::base64Decode(const std::string& encoded) {
    BIO* bio = BIO_new_mem_buf(encoded.data(), static_cast<int>(encoded.size()));
    BIO* b64 = BIO_new(BIO_f_base64());
    bio = BIO_push(b64, bio);
    BIO_set_flags(bio, BIO_FLAGS_BASE64_NO_NL);

    std::vector<unsigned char> result(encoded.size());
    int decodedLength = BIO_read(bio, result.data(), static_cast<int>(result.size()));
    BIO_free_all(bio);

    if (decodedLength < 0) {
        return {};
    }
    result.resize(decodedLength);
    return result;
}

// ==================== 随机数生成 ====================

std::string EncryptionHelper::generateKey() {
    std::vector<unsigned char> key(AES_KEY_SIZE);
    if (RAND_bytes(key.data(), AES_KEY_SIZE) != 1) {
        throw std::runtime_error("Failed to generate random key");
    }

    // 转换为Hex字符串
    std::ostringstream oss;
    for (unsigned char byte : key) {
        oss << std::hex << std::setw(2) << std::setfill('0') << static_cast<int>(byte);
    }
    return oss.str();
}

std::vector<unsigned char> EncryptionHelper::generateIV() {
    std::vector<unsigned char> iv(AES_IV_SIZE);
    if (RAND_bytes(iv.data(), AES_IV_SIZE) != 1) {
        throw std::runtime_error("Failed to generate random IV");
    }
    return iv;
}

// ==================== AES-256-GCM加密 ====================

std::vector<unsigned char> EncryptionHelper::aesGcmEncrypt(
    const std::string& plaintext,
    const std::vector<unsigned char>& key,
    const std::vector<unsigned char>& iv,
    std::vector<unsigned char>& tag) {

    EVP_CIPHER_CTX* ctx = EVP_CIPHER_CTX_new();
    if (!ctx) {
        throw std::runtime_error("Failed to create cipher context");
    }

    // 初始化加密操作
    if (EVP_EncryptInit_ex(ctx, EVP_aes_256_gcm(), nullptr, nullptr, nullptr) != 1) {
        EVP_CIPHER_CTX_free(ctx);
        throw std::runtime_error("Failed to initialize encryption");
    }

    // 设置IV长度
    if (EVP_CIPHER_CTX_ctrl(ctx, EVP_CTRL_GCM_SET_IVLEN, AES_IV_SIZE, nullptr) != 1) {
        EVP_CIPHER_CTX_free(ctx);
        throw std::runtime_error("Failed to set IV length");
    }

    // 设置密钥和IV
    if (EVP_EncryptInit_ex(ctx, nullptr, nullptr, key.data(), iv.data()) != 1) {
        EVP_CIPHER_CTX_free(ctx);
        throw std::runtime_error("Failed to set key and IV");
    }

    // 加密
    std::vector<unsigned char> ciphertext(plaintext.size() + EVP_CIPHER_block_size(EVP_aes_256_gcm()));
    int len = 0;
    if (EVP_EncryptUpdate(ctx,
                          ciphertext.data(), &len,
                          reinterpret_cast<const unsigned char*>(plaintext.c_str()),
                          static_cast<int>(plaintext.size())) != 1) {
        EVP_CIPHER_CTX_free(ctx);
        throw std::runtime_error("Encryption failed");
    }
    int ciphertextLen = len;

    // 完成加密
    if (EVP_EncryptFinal_ex(ctx, ciphertext.data() + len, &len) != 1) {
        EVP_CIPHER_CTX_free(ctx);
        throw std::runtime_error("Encryption finalization failed");
    }
    ciphertextLen += len;
    ciphertext.resize(ciphertextLen);

    // 获取认证标签
    tag.resize(AES_TAG_SIZE);
    if (EVP_CIPHER_CTX_ctrl(ctx, EVP_CTRL_GCM_GET_TAG, AES_TAG_SIZE, tag.data()) != 1) {
        EVP_CIPHER_CTX_free(ctx);
        throw std::runtime_error("Failed to get authentication tag");
    }

    EVP_CIPHER_CTX_free(ctx);
    return ciphertext;
}

std::string EncryptionHelper::aesGcmDecrypt(
    const std::vector<unsigned char>& ciphertext,
    const std::vector<unsigned char>& key,
    const std::vector<unsigned char>& iv,
    const std::vector<unsigned char>& tag) {

    EVP_CIPHER_CTX* ctx = EVP_CIPHER_CTX_new();
    if (!ctx) {
        return "";  // 解密失败
    }

    // 初始化解密操作
    if (EVP_DecryptInit_ex(ctx, EVP_aes_256_gcm(), nullptr, nullptr, nullptr) != 1) {
        EVP_CIPHER_CTX_free(ctx);
        return "";
    }

    // 设置IV长度
    if (EVP_CIPHER_CTX_ctrl(ctx, EVP_CTRL_GCM_SET_IVLEN, AES_IV_SIZE, nullptr) != 1) {
        EVP_CIPHER_CTX_free(ctx);
        return "";
    }

    // 设置密钥和IV
    if (EVP_DecryptInit_ex(ctx, nullptr, nullptr, key.data(), iv.data()) != 1) {
        EVP_CIPHER_CTX_free(ctx);
        return "";
    }

    // 解密
    std::vector<unsigned char> plaintext(ciphertext.size() + EVP_CIPHER_block_size(EVP_aes_256_gcm()));
    int len = 0;
    if (EVP_DecryptUpdate(ctx, plaintext.data(), &len,
                          ciphertext.data(), static_cast<int>(ciphertext.size())) != 1) {
        EVP_CIPHER_CTX_free(ctx);
        return "";
    }
    int plaintextLen = len;

    // 设置认证标签
    if (EVP_CIPHER_CTX_ctrl(ctx, EVP_CTRL_GCM_SET_TAG, AES_TAG_SIZE,
                            const_cast<unsigned char*>(tag.data())) != 1) {
        EVP_CIPHER_CTX_free(ctx);
        return "";
    }

    // 完成解密（会验证认证标签）
    if (EVP_DecryptFinal_ex(ctx, plaintext.data() + len, &len) != 1) {
        // 认证失败或数据被篡改
        EVP_CIPHER_CTX_free(ctx);
        return "";
    }
    plaintextLen += len;
    plaintext.resize(plaintextLen);

    EVP_CIPHER_CTX_free(ctx);
    return std::string(plaintext.begin(), plaintext.end());
}

// ==================== 公共接口 ====================

std::string EncryptionHelper::encrypt(const std::string& plaintext, const std::string& key) {
    // 1. 解析密钥（Hex -> bytes）
    if (key.size() != AES_KEY_SIZE * 2) {  // Hex编码是原始长度的2倍
        throw std::invalid_argument("Invalid key length");
    }

    std::vector<unsigned char> keyBytes(AES_KEY_SIZE);
    for (int i = 0; i < AES_KEY_SIZE; ++i) {
        std::string byteString = key.substr(i * 2, 2);
        keyBytes[i] = static_cast<unsigned char>(std::stoi(byteString, nullptr, 16));
    }

    // 2. 生成随机IV
    std::vector<unsigned char> iv = generateIV();

    // 3. 加密
    std::vector<unsigned char> tag;
    std::vector<unsigned char> ciphertext = aesGcmEncrypt(plaintext, keyBytes, iv, tag);

    // 4. 组合输出：IV + Ciphertext + Tag
    std::vector<unsigned char> output;
    output.insert(output.end(), iv.begin(), iv.end());
    output.insert(output.end(), ciphertext.begin(), ciphertext.end());
    output.insert(output.end(), tag.begin(), tag.end());

    // 5. Base64编码
    return base64Encode(output);
}

std::string EncryptionHelper::decrypt(const std::string& ciphertext, const std::string& key) {
    try {
        // 1. 解析密钥
        if (key.size() != AES_KEY_SIZE * 2) {
            return "";
        }

        std::vector<unsigned char> keyBytes(AES_KEY_SIZE);
        for (int i = 0; i < AES_KEY_SIZE; ++i) {
            std::string byteString = key.substr(i * 2, 2);
            keyBytes[i] = static_cast<unsigned char>(std::stoi(byteString, nullptr, 16));
        }

        // 2. Base64解码
        std::vector<unsigned char> data = base64Decode(ciphertext);
        if (data.size() < AES_IV_SIZE + AES_TAG_SIZE) {
            return "";  // 数据太短
        }

        // 3. 分离IV、Ciphertext、Tag
        std::vector<unsigned char> iv(data.begin(), data.begin() + AES_IV_SIZE);
        std::vector<unsigned char> tag(data.end() - AES_TAG_SIZE, data.end());
        std::vector<unsigned char> ciphertextBytes(
            data.begin() + AES_IV_SIZE,
            data.end() - AES_TAG_SIZE);

        // 4. 解密
        return aesGcmDecrypt(ciphertextBytes, keyBytes, iv, tag);

    } catch (...) {
        return "";  // 解密失败
    }
}

} // namespace MessageFramework
